import time
import math
import serial
import struct
from typing import Optional
from PyQt6.QtCore import QThread, QObject, QMutex, pyqtSignal

class UnitreeMotorProtocol:
    """
    Unitree电机通信协议处理类
    发送指令：17字节，帧头0xfeee
    接收反馈：15字节，帧头0xfdee
    最后两个字节都是CRC16_CCITT校验码
    """
    
    # 帧头定义
    TX_HEADER = bytes([0xfe, 0xee])  # 发送指令帧头
    RX_HEADER = bytes([0xfd, 0xee])  # 接收反馈帧头
    
    # 数据长度定义
    TX_PACKET_LENGTH = 17  # 发送指令总长度
    RX_PACKET_LENGTH = 16  # 接收反馈总长度
    
    def __init__(self):
        self.buffer = bytearray()

    def crc16_ccitt(self, data):
        """
        计算CRC16-CCITT (Kermit)校验码
        多项式：0x8408 (反向的0x1021)
        初始值：0x0000
        """
        crc = 0x0000
        polynomial = 0x8408
        
        for byte in data:
            crc ^= byte
            for _ in range(8):
                if crc & 0x0001:
                    crc = (crc >> 1) ^ polynomial
                else:
                    crc = crc >> 1
        
        return crc

    def data_received(self, data):
        """
        缓存接收到的数据，查找完整的数据包并处理
        """
        self.buffer.extend(data)
        results = []

        while self.RX_HEADER in self.buffer:
            header_pos = self.buffer.find(self.RX_HEADER)
            if (header_pos + self.RX_PACKET_LENGTH <= len(self.buffer)):
                packet = bytearray(self.buffer[header_pos:header_pos+self.RX_PACKET_LENGTH])
                self.buffer = bytearray(self.buffer[header_pos+self.RX_PACKET_LENGTH:])
                results.append(self.handle_rx_packet(packet))
                print(f"处理数据包: {packet.hex().upper()}")
            else:
                # 如果没有足够的数据包长度，退出循环等待更多数据
                break
        
        return results

    def handle_rx_packet(self, packet):
        """
        处理接收到的反馈数据包
        """
        result = {
            'identifier': '0x66',
            'data': {},
            'error': '',
        }
        
        try:
            # 验证帧头
            if packet[:2] != self.RX_HEADER:
                result['error'] = f"无效的帧头: {packet[:2].hex().upper()}, 期望: {self.RX_HEADER.hex().upper()}"
                return result
            
            # 提取数据部分（去掉帧头2字节和CRC2字节）
            data = packet[2:-2]
            received_crc = struct.unpack('<H', packet[-2:])[0]  # 小端序读取CRC
            
            # 计算CRC校验
            calculated_crc = self.crc16_ccitt(packet[:-2])
            
            if calculated_crc != received_crc:
                result['error'] = f"CRC校验失败: 计算值 0x{calculated_crc:04X}, 接收值 0x{received_crc:04X}"
                return result
            
            # 解包数据
            if len(data) != 12:
                return {'error': f"数据长度错误，期望12字节，实际{len(data)}字节"}
        
            # 1. 模式字节（ID和STATUS）
            mode_byte = data[0]
            motor_id = (mode_byte >> 4) & 0x0F  # 高4位为ID
            status = (mode_byte >> 1) & 0x07     # 中间3位为STATUS
            
            # 2. 实际输出转矩（τ_fbk）
            tau_fbk = struct.unpack('<h', data[1:3])[0] / 256.0  # 2字节有符号整数，除以256
        
            # 3. 实际角速度（ω_fbk）
            omega_fbk = struct.unpack('<h', data[3:5])[0] 
            omega_fbk = omega_fbk * (2 * math.pi) / 256.0  # 转换公式：ω(rad/s) = ω_fbk * 2π / 256
            
            # 4. 实际角度位置（θ_fbk）
            theta_fbk = struct.unpack('<i', data[5:9])[0]
            theta_fbk = theta_fbk * (2 * math.pi) / 32768.0  # 转换公式：θ(rad) = θ_fbk * 2π / 32768
            
            # 5. 电机温度（TEMP）
            temp = struct.unpack('<b', bytes([data[9]]))[0]  # 1字节有符号整数

            # 6. 电机错误标识
            motor_error = data[10] & 0x07  # 错误标识
            if motor_error > 0:
                result['error'] = f"电机错误标识: {motor_error}"
                return result
            
            # 7. 足端力（FORCE，12位）
            force_byte1 = data[10]
            force_byte2 = data[11]
            force_high = (force_byte1 & 0x1F)  # 低5位
            force_low = (force_byte2 & 0xFE) >> 1  # 高7位（去除保留位）
            force = (force_high << 7) | force_low
        
            result['data'] =  {
                'motor_id': motor_id,
                'status': status,
                'torque': tau_fbk,          # 单位：N·m
                'velocity': omega_fbk,      # 单位：rad/s
                'position': theta_fbk,      # 单位：rad
                'temperature': temp,        # 单位：℃
                'foot_force': force,        # 12位原始值
                'pc_timestamp': time.time_ns() / 1000000  # 添加PC时间戳
            }
            
        except Exception as e:
            result['error'] = f"数据包处理异常: {str(e)}"
        
        return result

    def pack_motor_command(self, command_data):
        """
        打包电机控制指令（17字节）
        格式参考手册表1：主机侧控制协议
        command_data需包含：id, mode, T, W, Pos, K_P, K_W
        """
        # 1. 帧头
        packet = bytearray()
        packet.extend(self.TX_HEADER)
        
        # 2. 模式字节（ID[4bit] + STATUS[3bit] + 保留[1bit]）
        motor_id = command_data.get('id', 0)
        mode = command_data.get('mode', 0)  # 0:锁定, 1:FOC, 2:校准
        if motor_id < 0 or motor_id > 15:
            raise ValueError(f"电机ID必须为0-15，实际为{motor_id}")
        if mode < 0 or mode > 7:
            raise ValueError(f"模式必须为0-7，实际为{mode}")
        mode_byte =  motor_id | (mode << 4)  # 保留位设为0
        packet.append(mode_byte)
        
        # 3. 前馈力矩（τ_set = T * 256，2字节有符号整数）
        t_ff = command_data.get('T', 0.0)
        if abs(t_ff) > 127.99:
            raise ValueError(f"前馈力矩超出范围[-127.99, 127.99]，实际为{t_ff}")
        tau_set = int(round(t_ff * 256))
        packet.extend(struct.pack('<h', tau_set))
        
        # 4. 期望角速度（ω_set = W * 128，2字节有符号整数）
        w_des = command_data.get('W', 0.0)
        if abs(w_des) > 804.0:
            raise ValueError(f"期望角速度超出范围[-804.0, 804.0]，实际为{w_des}")
        omega_set = int(round(w_des * 128))
        packet.extend(struct.pack('<h', omega_set))
        
        # 5. 期望角度位置（θ_set = Pos * 32768 / (2π)，4字节有符号整数）
        pos_des = command_data.get('Pos', 0.0)
        theta_set = int(round(pos_des * 32768 / (2 * math.pi)))
        packet.extend(struct.pack('<i', theta_set))
        
        # 6. 位置刚度（K_pos = K_P * 1280，2字节无符号整数）
        k_p = command_data.get('K_P', 0.0)
        if k_p < 0 or k_p > 25.599:
            raise ValueError(f"位置刚度超出范围[0, 25.599]，实际为{k_p}")
        k_pos = int(round(k_p * 1280))
        packet.extend(struct.pack('<H', k_pos))
        
        # 7. 速度刚度（K_spd = K_W * 1280，2字节无符号整数）
        k_w = command_data.get('K_W', 0.0)
        if k_w < 0 or k_w > 25.599:
            raise ValueError(f"速度刚度超出范围[0, 25.599]，实际为{k_w}")
        k_spd = int(round(k_w * 1280))
        packet.extend(struct.pack('<H', k_spd))
        
        # 8. CRC校验
        crc = self.crc16_ccitt(packet)
        packet.extend(struct.pack('<H', crc))
        
        if len(packet) != self.TX_PACKET_LENGTH:
            raise RuntimeError(f"指令长度错误，期望17字节，实际{len(packet)}字节")
        
        return bytes(packet)


class UnitreeMotorThread(QThread):
    """
    Unitree电机通信线程
    """
    RX_READ_LENGTH = 15
    data_received_signal = pyqtSignal(list)

    def __init__(self, serial_instance):
        super().__init__()
        self.serial = serial_instance
        self.protocol = UnitreeMotorProtocol();
        self._running = True
        self.mutex = QMutex()

    def run(self):
        """线程运行主循环"""
        if not hasattr(self.serial, 'cancel_read'):
            self.serial.timeout = 1

        while self._running and self.serial.is_open:
            try:
                data = self.serial.read(self.RX_READ_LENGTH)
            except serial.SerialException as e:
                self._running = False
                break
            else:
                self.mutex.lock()
                results = self.protocol.data_received(data)
                self.mutex.unlock()
                if len(results) > 0:
                    self.data_received_signal.emit(results)


    def write(self, command_data):
        """
        发送控制指令到电机
        
        Args:
            command_data: 控制指令数据字典
            
        Returns:
            str: 发送结果消息
        """

        try:
            self.mutex.lock()
            packet = self.protocol.pack_motor_command(command_data)
            self.mutex.unlock()
            self.serial.write(packet)
            result = True
        except Exception as e:
            print(f"指令打包或发送失败: {str(e)}")
            result = False
        
        return result

    def close(self):
        """关闭线程"""
        self._running = False
        if hasattr(self.serial, 'cancel_read'):
            self.serial.cancel_read()
        self.wait(100)
        self.serial.close()


class UnitreeMotorService(QObject):
    """
    Unitree电机通信服务
    """
    
    message_signal = pyqtSignal(str)

    def __init__(self, callback):
        super().__init__()
        self.ser: Optional[serial.Serial] = None
        self.port: Optional[str] = None
        self.baudrate: Optional[int] = None
        self.protocol = UnitreeMotorProtocol()
        self.motor_thread: Optional[UnitreeMotorThread] = None
        self.callback = callback

    def open_serial(self, port: str, baudrate: int = 4000000):
        """
        打开串口并启动通信线程
        
        Args:
            port: 串口端口号
            baudrate: 波特率，默认4000000（Unitree电机常用波特率）
            
        Returns:
            bool: 成功返回True，失败返回False
        """
        self.port = port
        self.baudrate = baudrate
        try:
            self.ser = serial.Serial(port, baudrate)
            self.motor_thread = UnitreeMotorThread(self.ser)
            self.motor_thread.data_received_signal.connect(self.callback)
            self.motor_thread.start()
            self.message_signal.emit(f"Unitree电机串口打开成功: {port}@{baudrate}")
            return True
        except Exception as e:
            self.message_signal.emit(f"Unitree电机串口打开失败: {e}")
            self.ser = None
            return False

    def close_serial(self):
        """关闭串口和通信线程"""
        if self.motor_thread:
            self.motor_thread.close()
            self.motor_thread = None
            self.ser = None
            self.message_signal.emit("Unitree电机串口已关闭")
        else:
            self.message_signal.emit("Unitree电机串口未打开，无法关闭")

    def send_command(self, command_data):
        """
        发送控制指令到电机
        
        Args:
            command_data: 控制指令数据字典
        
        Note:
            command_data = {
                'id': 0,        # 电机ID，0-15
                'mode': 1,      # 控制模式，0:锁定, 1:FOC, 2:校准
                'T': 0.0,       # 前馈力矩，单位N·m，范围[-127.99, 127.99]
                'W': 0.0,       # 期望角速度，单位rad/s，范围[-804.0, 804.0]
                'Pos': 0.0,     # 期望角度位置，单位rad
                'K_P': 0.0,     # 位置刚度，范围[0, 25.599]
                'K_W': 0.0,     # 速度刚度，范围[0, 25.599]
            }
            
        Returns:
            bool: 成功返回True，失败返回False
        """
        if not (self.ser and self.ser.is_open):
            self.message_signal.emit("Unitree电机串口未打开，无法发送指令")
            return False
        
        if not (self.motor_thread and self.motor_thread.isRunning()):
            self.message_signal.emit("Unitree电机通信线程未运行，无法发送指令")
            return False
            
        result = self.motor_thread.write(command_data)
        return result

    def is_connected(self):
        """
        检查连接状态
        
        Returns:
            bool: 已连接返回True，未连接返回False
        """
        return self.ser is not None and self.ser.is_open
